import math
from typing import Dict, List
import random
from utils.Seed import Seed
from schedule.PowerSchedule import PowerSchedule


class FairSchedule(PowerSchedule):
    def __init__(self, threshold: int = 0) -> None:
        super().__init__()
        self.threshold = threshold
        self.path_stats: Dict[int, int] = (
            None  # 路径统计 {path_id: {'freq': int, 'selections': int}}
        )

    def bind_frequency(self, freq_dict: Dict[int, int]) -> None:
        self.path_stats = freq_dict

    def normalized_energy(self, population: List[Seed]) -> List[float]:
        """Normalize energy"""
        energy = list(map(lambda seed: seed.energy, population))
        sum_energy = sum(energy)  # Add up all values in energy
        if sum_energy == 0:  # 可能发生
            sum_energy = len(energy)
            norm_energy = list(map(lambda nrg: 1 / sum_energy, energy))
        else:
            norm_energy = list(map(lambda nrg: nrg / sum_energy, energy))
        return norm_energy

    def assign_energy(self, population: List[Seed]) -> None:
        if not self.path_stats:
            return
        else:
            """
            动态计算阈值
            """
            if self.threshold == 0:
                freqs = self.path_stats.values()
                min_freqs = min(freqs)
                threshold = 2 ** int(math.log2(min_freqs) + 1)
            else:
                threshold = self.threshold

            for seed in population:
                path_id = hash(frozenset(seed.coverage))
                freq = self.path_stats.get(path_id, 0)
                if freq > threshold:
                    # 直接抛弃高频路径
                    seed.energy = 0
                else:
                    seed.energy = threshold / (freq + 1)
